#include <iostream>
#include <fstream>
#include <cmath>
#include <vector>
#include <cstdlib>
#include <cstdio>
#include "img_arith.h"
#include "sample.h"
#include "image.h"
#include "utils/device_property.h"
#include "math/math_ops.h"
#include "math/reduction.h"
#include "utils/assertion.h"
#include "utils/data_filler.h"
#include "tensor.h"
#include "layers/conv1d.h"
#include "layers/conv2d.h"

#define N 10000

void test_image_arithmetic()
{
    std::string image_path("../data/person.jpg");
    std::string image_path2("../data/person2.jpg");
    uzu::ImageU8 image;
    uzu::ReadImage(image_path, image);

    uzu::ImageU8 gray;
    uzu::rgb2gray(image, gray);

    // blur image
    uzu::ImageU8 blurred;
    uzu::boxBlur(gray, blurred, 7);

    uzu::WriteImage(image_path2, blurred);
}

void test_device()
{
    uzu::DeviceProperty property;
    std::cout << property.NumDevices() << std::endl;
    std::cout << property.MaxThreadsPerBlock(0) << std::endl;
}

void test_image_io()
{
    std::string image_path("../data/person.jpg");
    std::string image_path2("../data/person2.jpg");
    uzu::Image<uint8_t> image;
    uzu::ReadImage(image_path, image);
    uzu::WriteImage(image_path2, image);
}

void test_matrix_mul()
{
    uzu::MatrixF mat1(1000, 256);
    uzu::MatrixF mat2(256, 324);
    uzu::FillRandom(mat1);
    uzu::FillRandom(mat2);

    uzu::MatrixF res1 = mat1 * mat2;
    for (int i = 0; i < 10; ++i)
        std::cout << res1.data[i] << std::endl;
    
    uzu::MatrixF res2 = uzu::MatrixMul(mat1, mat2);
    for (int i = 0; i < 10; ++i)
        std::cout << res2.data[i] << std::endl;

    if (uzu::Assertion::Equal(res1.data, res2.data, res1.rows * res1.cols))
        std::cout << "Equal\n";
    else
        std::cout << "Not Equal\n";
}


void test_reduce_sum()
{
    int size = 1000;
    float* data = new float[size];
    uzu::FillRandom(data, size);
    float s = uzu::ReduceSumCpu(data, size);
    std::cout << s << std::endl;
    float s2 = uzu::ReduceSum(data, size);
    std::cout << s2 << std::endl;
}


void test_conv1d()
{
    const int data_size = 2000;
    std::vector<float> vsource;
    for (int i = 0; i < data_size; ++i)
        vsource.push_back(i * 1.0f);
    
    std::vector<float> vfilter = {1, 1, 1, 1, 1};
    uzu::Tensor source(data_size);
    uzu::Tensor filter(5);
    source.CopyFromHost(vsource.data());
    filter.CopyFromHost(vfilter.data());
    uzu::Tensor dest;
    uzu::Conv1d(source, filter, dest);
    dest.ToHost();
    // test
    uint32_t size = dest.Size();
    float* data = dest.Data();
    for (uint32_t i = 0; i < size; ++i)
    {
        std::cout << data[i] << std::endl;
    }
}


void LoadFromTextFile(const std::string& filePath, std::vector<float>& data)
{
    std::ifstream in;
    in.open(filePath);
    if (!in.is_open()) return;

    std::string line;
    while (true)
    {
        std::getline(in, line);
        if (line.empty()) break;
        float value = std::atof(line.c_str());
        data.push_back(value);
    }
}

void show_tensor(uzu::Tensor& tensor, int num_elements)
{
    tensor.ToHost();
    float* data = tensor.Data();
    uint32_t size = tensor.Size();
    int n = std::min(int(size), num_elements);
    for (int i = 0; i < n; ++i)
        std::cout << data[i] << std::endl;
}

void write_tensor(uzu::Tensor& tensor, const std::string& filepath)
{
    FILE* fp;
    fp = fopen(filepath.c_str(), "w");
    uint32_t size = tensor.Size();
    tensor.ToHost();
    float* data = tensor.Data();
    for (uint32_t i = 0; i < size; ++i)
    {
        fprintf(fp, "%.04f\n", data[i]);
    }
    fclose(fp);
}

void test_conv2d()
{
    uzu::Tensor source(100, 100);
    std::vector<float> vsource;
    LoadFromTextFile("../data/sample_input.txt", vsource);
    source.CopyFromHost(vsource.data());
    uzu::Tensor filter(5, 5);
    std::vector<float> vfilter;
    LoadFromTextFile("../data/sample_filter.txt", vfilter);
    filter.CopyFromHost(vfilter.data());
    uzu::Tensor sample_output(100, 100);
    std::vector<float> vout;
    LoadFromTextFile("../data/sample_output.txt", vout);
    sample_output.CopyFromHost(vout.data());
    show_tensor(sample_output, 10);
    // 
    uzu::Tensor dest;
    uzu::Conv2d(source, filter, dest);
    dest.ToHost();
    show_tensor(dest, 10);
}


void test_reduction()
{
    int num = 1000000;
    std::vector<float> varray;
    for (int i = 0; i < num; ++i)
        varray.push_back(1);
    
    uzu::Tensor inputs(num);
    inputs.CopyFromHost(varray.data());

    uzu::Tensor outputs;
    float result = uzu::ReduceSum(inputs);
    std::cout << result << std::endl;
}

void test_histogram()
{
    std::string image_path("../data/person.jpg");
    uzu::ImageU8 image;
    uzu::ReadImage(image_path, image);
    uzu::ImageU8 gray;
    uzu::rgb2gray(image, gray);

    gray.ToHost();
    uzu::Tensor hist_cpu;
    uzu::histogramCpu(gray, hist_cpu);
    gray.ToDevice();
    uzu::Tensor hist_gpu;
    uzu::histogram(gray, hist_gpu);

    show_tensor(hist_cpu, 10);
    show_tensor(hist_gpu, 10);
}


int main()
{
    // test_device();
    // test_reduce_sum();
    // test_conv1d();
    // test_conv2d();
    // test_reduction();
    test_histogram();

    return 0;
}
